import open3d as o3d
import numpy as np
import os
from pathlib import Path

def viewpoint_to_transform_matrix(viewpoint_str):
    """
    将PCD文件头中的VIEWPOINT字符串转换为4x4变换矩阵。
    格式: "VIEWPOINT tx ty tz qw qx qy qz"
    """
    # 分割字符串并提取7个数值，转换为浮点数
    parts = viewpoint_str.split()
    if len(parts) != 8: # 包括"VIEWPOINT"这个词本身，所以是8部分
        raise ValueError(f"VIEWPOINT format error. Expected 8 parts, got {len(parts)}: {viewpoint_str}")
    
    # 提取平移和四元数
    tx, ty, tz = float(parts[1]), float(parts[2]), float(parts[3])
    qw, qx, qy, qz = float(parts[4]), float(parts[5]), float(parts[6]), float(parts[7])
    
    # 根据四元数计算旋转矩阵
    # 公式: R = [[1-2*(qy^2+qz^2), 2*(qx*qy - qw*qz),   2*(qx*qz + qw*qy)],
    #            [2*(qx*qy + qw*qz),   1-2*(qx^2+qz^2), 2*(qy*qz - qw*qx)],
    #            [2*(qx*qz - qw*qy),   2*(qy*qz + qw*qx),   1-2*(qx^2+qy^2)]]
    R = np.array([
        [1 - 2*(qy*qy + qz*qz),      2*(qx*qy - qw*qz),      2*(qx*qz + qw*qy)],
        [2*(qx*qy + qw*qz),      1 - 2*(qx*qx + qz*qz),      2*(qy*qz - qw*qx)],
        [2*(qx*qz - qw*qy),      2*(qy*qz + qw*qx),      1 - 2*(qx*qx + qy*qy)]
    ])
    
    # 构建4x4齐次变换矩阵
    transform_matrix = np.identity(4)
    transform_matrix[:3, :3] = R  # 设置旋转部分
    transform_matrix[:3, 3] = [tx, ty, tz] # 设置平移部分
    
    return transform_matrix

def parse_viewpoint_from_pcd_header(file_path):
    """
    读取PCD文件头，找到并返回VIEWPOINT行。
    """
    viewpoint_line = None
    with open(file_path, 'r') as f:
        # 读取文件头行，直到遇到"DATA"
        for line in f:
            if line.startswith('VIEWPOINT'):
                viewpoint_line = line.strip()
            if line.startswith('DATA'):
                # 读到DATA行说明文件头结束
                break
    if viewpoint_line is None:
        raise ValueError(f"VIEWPOINT not found in the header of {file_path}")
    return viewpoint_line

def merge_pcds_with_viewpoint(pcd_dir, output_file="merged_pcb.pcd", voxel_size=0.0):
    """
    主函数：合并一个目录下所有具有VIEWPOINT信息的PCD文件。
    
    参数:
        pcd_dir: 包含PCD文件的目录路径
        output_file: 输出文件名
        voxel_size: 体素下采样大小。合并后点云可能很大，可用此参数简化。
                    >0时生效，例如0.01表示1cm网格下采样。0表示不下采样。
    """
    # 获取目录下所有.pcd文件
    pcd_paths = sorted(Path(pcd_dir).glob('*.pcd'))
    if not pcd_paths:
        print(f"No PCD files found in {pcd_dir}")
        return
    
    print(f"Found {len(pcd_paths)} PCD files.")
    
    transformed_clouds = []
    
    for pcd_path in pcd_paths:
        try:
            print(f"Processing {pcd_path.name}...")
            
            # 1. 从文件头解析VIEWPOINT
            viewpoint_str = parse_viewpoint_from_pcd_header(pcd_path)
            print(f"  Viewpoint: {viewpoint_str}")
            
            # 2. 将VIEWPOINT转换为变换矩阵
            trans_matrix = viewpoint_to_transform_matrix(viewpoint_str)
            
            # 3. 用Open3D读取点云
            # 注意：必须用`read_point_cloud`直接读，我们自己解析的header只用于位姿
            cloud = o3d.io.read_point_cloud(str(pcd_path))
            
            # 4. 应用变换矩阵
            cloud.transform(trans_matrix)
            transformed_clouds.append(cloud)
            print(f"  Successfully transformed.\n")
            
        except Exception as e:
            print(f"  Error processing {pcd_path.name}: {e}. Skipping.\n")
    
    # 合并所有点云
    if not transformed_clouds:
        print("No point clouds were successfully processed.")
        return
    
    print("Merging all point clouds...")
    # 高效合并方式：先将所有点云放入一个列表，然后一起合并
    merged_cloud = transformed_clouds[0]
    for i in range(1, len(transformed_clouds)):
        merged_cloud += transformed_clouds[i]
    
    # (可选) 体素下采样以减少数据量
    if voxel_size > 0:
        print(f"Downsampling with voxel size {voxel_size}...")
        merged_cloud = merged_cloud.voxel_down_sample(voxel_size)
    
    # 可视化结果
    print("Displaying merged point cloud (close window to continue)...")
    o3d.visualization.draw_geometries([merged_cloud])
    
    # 保存结果
    print(f"Saving merged point cloud to {output_file}...")
    # 保存为PCD或PLY格式
    if output_file.endswith('.pcd'):
        o3d.io.write_point_cloud(output_file, merged_cloud)
    else:
        # PLY格式更通用，文件更小
        o3d.io.write_point_cloud(output_file, merged_cloud)
    print("All done!")

# 使用示例
if __name__ == "__main__":
    # 设置您的PCD文件目录路径
    pcd_directory = "/path/to/your/pcd/files"  # 请修改为您的实际路径
    
    # 调用函数进行合并
    merge_pcds_with_viewpoint(
        pcd_dir=pcd_directory,
        output_file="merged_pcb_project.ply", # 输出为PLY文件，也可以改为 .pcd
        voxel_size=0.002  # 可选：2mm下采样，根据你的点云密度调整，如果不需要设为0
    )